import torch  # 导入PyTorch深度学习框架
from torch import nn  # 从torch导入神经网络模块

import cv2  # 导入OpenCV库，用于图像处理
from torchvision import transforms  # 导入torchvision的transforms模块，用于图像预处理
from ResNet import resnet18  # 从ResNet模块导入resnet18模型定义

# 定义类别字典，将模型输出的数字标签映射为实际类别名称
classes = {0: "half-ripe", 1: "raw", 2: "ripe"}

# 1、加载模型
# 假设模型定义在`resnet18`函数中并已保存最优模型为'./model/best.pth'
model = resnet18().to('cpu')  # 初始化resnet18模型并将其移动到CPU（根据实际模型定义调整）

# 获取全连接层的输入特征数
fc_inputs = model.fc.in_features
# 重新定义全连接层，将输出维度改为3（对应3个类别）
model.fc = nn.Sequential(
    nn.Linear(fc_inputs, 3),  # 线性层：输入为fc_inputs，输出为3
).to('cpu')  # 将新定义的层移动到CPU

# 加载预训练好的模型权重
model.load_state_dict(torch.load('../model/best.pth', map_location=torch.device('cpu')))

model.eval()  # 设置模型为评估模式（关闭 dropout、批量归一化使用移动平均值等）

# 2、模型预热，确保在不计算梯度的情况下进行推断
with torch.no_grad():  # 上下文管理器，禁用梯度计算，节省内存并加速计算
    # 输入一个全零的张量（1个样本，3个通道，224x224大小）进行一次前向传播
    model(torch.zeros(size=(1, 3, 224, 224)).to('cpu'))

# 3、推理相关设置
# 定义图像预处理步骤，torchvision.transforms通常需要PIL图像对象作为输入
preprocess = transforms.Compose([
    transforms.ToPILImage(),  # 先将OpenCV的numpy数组转换为PIL图像
    transforms.Resize((224, 224)),  # 将图像大小调整为224x224
    transforms.ToTensor(),  # 将PIL图像转换为Tensor，并将像素值归一化到[0,1]
    # 对图像进行标准化，使用ImageNet数据集的均值和标准差
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])


def predict_resnet(image):
    """使用ResNet模型对输入图像进行预测，返回类别名称"""
    # 将OpenCV默认的BGR格式图像转换为RGB格式（与训练时的格式一致）
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 对图像进行预处理
    input_tensor = preprocess(image)
    # 增加一个批次维度（变为[1, 3, 224, 224]）并移动到CPU
    input_batch = input_tensor.unsqueeze(0).to('cpu')

    # 确保在不计算梯度的情况下进行推断
    with torch.no_grad():
        output = model(input_batch)  # 模型前向传播，得到原始输出
        # 对输出应用softmax函数，得到每个类别的概率
        probabilities = torch.nn.functional.softmax(output[0], dim=0)

    # 获取概率最大的类别索引
    _, predicted_class = torch.max(probabilities, dim=0)
    # 返回对应的类别名称
    return classes[predicted_class.item()]